
#include "KDTree.h"

KDTree::KDTree(int maxDepth,int boundingDepth,size_t minTriangles)
{
	m_maxDepth = maxDepth;
	m_boundingDepth = boundingDepth;
	m_minTriangles = minTriangles;

	m_RenderTimes = 0;
	m_pRoot = new KDNode();
}

KDTree::~KDTree()
{
	// delete m_vpTriangles
	for (size_t i=0;i<m_pRoot->vpTriangles.size();i++)
		delete (Triangle*)(m_pRoot->vpTriangles[i]);

	// delete KDTree
	deleteKDNode(m_pRoot);
}

void KDTree::addTriangle(Triangle *pTriangle)
{
	m_pRoot->addTriangle(pTriangle);
}

void KDTree::deleteKDNode(KDNode *pNode)
{
	if ( pNode->pLeftNode )
		deleteKDNode( pNode->pLeftNode );
	if ( pNode->pRightNode )
		deleteKDNode( pNode->pRightNode );
	delete pNode;
}

void KDTree::createTree()
{
	createNode(m_pRoot,0,0);
}

void KDTree::createNode(KDNode *pNode, int depth, int lastSplitXYZ)
{
	if (pNode->vpTriangles.size()>m_minTriangles && depth<m_maxDepth)
	{
		KDNode *pLeftNode = new KDNode;
		KDNode *pRightNode = new KDNode;

		float dx = pNode->aabb.maxXYZ[0] - pNode->aabb.minXYZ[0];
		float dy = pNode->aabb.maxXYZ[1] - pNode->aabb.minXYZ[1];
		float dz = pNode->aabb.maxXYZ[2] - pNode->aabb.minXYZ[2];
		
		float maxValue = dx;
		pNode->splitXYZ = SPLITX;
		if (dy > maxValue) {
			pNode->splitXYZ = SPLITY;
			maxValue = dy; }
		if (dz > maxValue) {
			pNode->splitXYZ = SPLITZ;
			maxValue = dz;
		}
		switch (pNode->splitXYZ){
			case SPLITX: pNode->splitValue = pNode->aabb.minXYZ[0] + maxValue/2.f; break;
			case SPLITY: pNode->splitValue = pNode->aabb.minXYZ[1] + maxValue/2.f; break;
			case SPLITZ: pNode->splitValue = pNode->aabb.minXYZ[2] + maxValue/2.f; break;
		}
// 		pNode->splitXYZ = lastSplitXYZ = (lastSplitXYZ+1)%3;
// 		switch (pNode->splitXYZ){
//  			case SPLITX: pNode->splitValue = (pNode->aabb.minXYZ[0] + pNode->aabb.maxXYZ[0])/2.f; break;
// 			case SPLITY: pNode->splitValue = (pNode->aabb.minXYZ[1] + pNode->aabb.maxXYZ[1])/2.f; break;
// 			case SPLITZ: pNode->splitValue = (pNode->aabb.minXYZ[2] + pNode->aabb.maxXYZ[2])/2.f; break;
// 		}
		for (size_t i=0; i<pNode->vpTriangles.size(); i++){
			Triangle *pTriangle = pNode->vpTriangles[i];
			float minV = pTriangle->m_AABB.minXYZ[pNode->splitXYZ];
			float maxV = pTriangle->m_AABB.maxXYZ[pNode->splitXYZ];
			if ( maxV <= pNode->splitValue - FLT_EPSILON )
				pLeftNode->addTriangle(pTriangle);
			else if ( minV >= pNode->splitValue + FLT_EPSILON )
				pRightNode->addTriangle(pTriangle);
			else {
				pLeftNode->addTriangle(pTriangle);
				pRightNode->addTriangle(pTriangle);
			}
		}
		createNode(pLeftNode,depth+1,lastSplitXYZ);
		createNode(pRightNode,depth+1,lastSplitXYZ);
		pNode->pLeftNode = pLeftNode;
		pNode->pRightNode = pRightNode;
	}
}

bool KDTree::searchNode(const KDNode *pNode, const Ray *pRay, 
				Hit *pHit, float tmin, float tmax, int depth)
{
 	float ttmin, ttmax;
	if ( depth < m_boundingDepth ){
	 	pNode->aabb.intersect(*pRay,ttmin,ttmax);
	 	ttmin = ttmin>tmin? ttmin:tmin;
	 	ttmax = ttmax<tmax? ttmax:tmax;
	 	if( ttmin > ttmax ) return false;
	} else {
		ttmin = tmin; ttmax = tmax;
	}

	if ( pNode->vpTriangles.size()==0 )
		return 0;
	if ( pNode->pLeftNode==NULL && pNode->pRightNode==NULL )
	{
		if ( pHit )	return pNode->intersect(pRay,pHit,ttmin,ttmax,m_RenderTimes);
		return pNode->intersectShadowRay(pRay,ttmin,ttmax,m_RenderTimes);
	}
	Vec3f R = pRay->getOrigin(); Vec3f D = pRay->getDirection();
	float thit;
	KDNode *pNear,*pFar;
// 	if ( D[pNode->splitXYZ] > -FLT_EPSILON && D[pNode->splitXYZ] < FLT_EPSILON )
// 	{
// 		if ( R[pNode->splitXYZ] >= pNode->splitValue + FLT_EPSILON )
// 			return searchNode(pNode->pRightNode,pRay,pHit,ttmin,ttmax,depth+1);
// 		else
// 			return searchNode(pNode->pLeftNode,pRay,pHit,ttmin,ttmax,depth+1);
// 	}
	thit = (pNode->splitValue-R[pNode->splitXYZ]) / D[pNode->splitXYZ];
	if ( D[pNode->splitXYZ] >= FLT_EPSILON ){
		pNear = pNode->pLeftNode;
		pFar = pNode->pRightNode;
	} else {
		pNear = pNode->pRightNode;
		pFar = pNode->pLeftNode;
	}
	
	if ( thit > ttmax )
		return searchNode(pNear,pRay,pHit,ttmin,ttmax,depth+1);
	if ( thit < ttmin )
		return searchNode(pFar,pRay,pHit,ttmin,ttmax,depth+1);
	if ( searchNode(pNear,pRay,pHit,ttmin,thit,depth+1) )
			return true;
	return searchNode(pFar,pRay,pHit,thit,ttmax,depth+1);
}

bool KDTree::intersect(const Ray *pRay, Hit *pHit, float tmin, float tmax)
{
	m_RenderTimes++;
	//return m_pRoot->intersect(pRay,pHit,tmin,tmax,m_RenderTimes);
 	return searchNode(m_pRoot,pRay,pHit,tmin,tmax,0);
}

bool KDTree::intersectShadowRay(const Ray *pRay, float tmin, float tmax)
{
	m_RenderTimes++;
	//return m_pRoot->intersectShadowRay(pRay,tmin,tmax,m_RenderTimes);
	return searchNode(m_pRoot,pRay,NULL,tmin,tmax,0);
	return false;
}

void KDTree::debugMessage()
{
	printf("KDTree{\n}\n\n");
}